import cv2
import mediapipe as mp
from pi_driver.msg import Point3D
from pi_driver.srv import GetLandmarks, GetLandmarksResponse

mp_drawing = mp.solutions.drawing_utils
mp_drawing_styles = mp.solutions.drawing_styles
mp_face_mesh = mp.solutions.face_mesh
drawing_spec = mp_drawing.DrawingSpec(thickness=1, circle_radius=1)

class FaceMesher:
    def __init__(self) -> None:
        self.face_mesh = mp_face_mesh.FaceMesh(
            max_num_faces=1,
            refine_landmarks=True,
            min_detection_confidence=0.5,
            min_tracking_confidence=0.5)

    def detect(self, image):
        h, w, channel = image.shape
        # To improve performance, optionally mark the image as not writeable to
        # pass by reference.
        # image.flags.writeable = False
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        results = self.face_mesh.process(image)
        # Draw the hand annotations on the image.
        # image.flags.writeable = True
        # image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        array = GetLandmarksResponse()
        if results.multi_face_landmarks:
            # print(results.multi_face_landmarks[0])
            for index, face_landmarks in enumerate(results.multi_face_landmarks):
                # print(face_landmarks.landmark, len(face_landmarks.landmark))
                array.landmarks = [Point3D(p.x,p.y,p.z) for p in face_landmarks.landmark]
                # array["keypoints"] = points
                    # "score": results.multi_handedness[index].classification[0].score,
                    # "label": results.multi_handedness[index].classification[0].label,
        return array, results

    def draw_results(self,image,results):
        if results.multi_face_landmarks:
            for face_landmarks in results.multi_face_landmarks:
                mp_drawing.draw_landmarks(
                    image=image,
                    landmark_list=face_landmarks,
                    connections=mp_face_mesh.FACEMESH_TESSELATION,
                    landmark_drawing_spec=None,
                    connection_drawing_spec=mp_drawing_styles
                    .get_default_face_mesh_tesselation_style())
                mp_drawing.draw_landmarks(
                    image=image,
                    landmark_list=face_landmarks,
                    connections=mp_face_mesh.FACEMESH_CONTOURS,
                    landmark_drawing_spec=None,
                    connection_drawing_spec=mp_drawing_styles
                    .get_default_face_mesh_contours_style())
                mp_drawing.draw_landmarks(
                    image=image,
                    landmark_list=face_landmarks,
                    connections=mp_face_mesh.FACEMESH_IRISES,
                    landmark_drawing_spec=None,
                    connection_drawing_spec=mp_drawing_styles
                    .get_default_face_mesh_iris_connections_style())
        # return image
if __name__ == '__main__':
    FaceMesher()
